# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: yuxuan

import datetime
import glob

import fastfeature.simulator_spec_pb2 as sspec
import google.protobuf.json_format as json_format

from coin.exchange.base.order_gateway import OrderSide
from coin.proto.coin_strategy_pb2 import StrategyLog
from coin.support.proto_log.logic.util import run_from_strat_log_archive


class OnLog(object):
  def on_submit(self, ts, currency, order_id, order_side, order_price, order_qty):
    pass

  def on_order_accept(self, ts, currency, order_id):
    pass

  def on_submit_reject(self, ts, currency, order_id, error_code, exchange_error_code, error_msg):
    pass

  def on_fill(self, ts, currency, order_id, order_side, fill_price, fill_qty, maker):
    pass

  def on_balance(self, ts, currency, balance):
    pass

  def on_position(self, ts, currency, position):
    pass

  def on_reserve(self, ts, currency, reserve):
    pass

  def on_cancel_submit(self, ts, currency, order_id):
    pass

  def on_cancel_accept(self, ts, currency, order_id):
    pass

  def on_cancel_confirm(self, ts, currency, order_id):
    pass

  def on_cancel_error(self, ts, currency, order_id, error_code, exchange_error_code, error_msg):
    pass

  def on_cancel_reject(self, ts, currency, order_id, error_code, exchange_error_code, error_msg):
    pass

  def on_config(self, ts, symbol, lot_size, min_pos, max_pos):
    pass


class _OnLog(object):
  def __init__(self, onlog):
    self._ts = 0
    if not isinstance(onlog, list):
      self._onlog = [onlog]
    else:
      self._onlog = onlog

  def _onfunc(self, onfunc, *argv):
    for onlog in self._onlog:
      getattr(onlog, onfunc)(*argv)

  def _on_reserve(self, reserve):
    for data in reserve:
      self._onfunc('on_reserve', self._ts, data['currency'], data['total'])

  def _on_submit(self, submit):
    order_id = submit.get('proc_order_id', submit.get('internal_order_id'))
    order_side = OrderSide.BUY if 'BUY' in submit['order_side'] else OrderSide.SELL
    order_price = submit['order_price']
    order_qty = submit['order_qty']
    self._onfunc('on_submit',
                 self._ts,
                 submit['symbol'],
                 order_id,
                 order_side,
                 order_price,
                 order_qty)

  def _on_submit_reject(self, submit_reject):
    symbol = submit_reject['symbol']
    order_id = submit_reject.get('proc_order_id', submit_reject.get('internal_order_id'))
    self._onfunc('on_submit_reject',
                 self._ts,
                 symbol,
                 order_id,
                 submit_reject.get('error_code'),
                 submit_reject.get('exchange_error_code'),
                 submit_reject.get('exchange_error_msg'))

  def _on_order_accept(self, ack):
    order_id = ack.get('proc_order_id', ack.get('internal_order_id'))
    self._onfunc('on_order_accept', self._ts, ack['symbol'], order_id)

  def _on_fill(self, fill):
    order_id = fill.get('proc_order_id', fill.get('internal_order_id'))
    symbol = fill['symbol']
    order_side = OrderSide.BUY if 'BUY' in fill['order_side'] else OrderSide.SELL
    fill_price = fill['fill_price']
    fill_qty = fill['fill_qty']
    maker = None
    fill_type = fill.get('fill_type', fill.get('guess_fill_type', None))
    if fill_type is not None:
      maker = 'MAKER' in fill_type
    self._onfunc('on_fill', self._ts, symbol, order_id, order_side, fill_price, fill_qty, maker)

  def _on_balance(self, balance):
    for abal in balance:
      self._onfunc('on_balance', self._ts, abal['currency'], abal['total'])

  def _on_position(self, position):
    for apos in position:
      if 'symbol' in apos and 'net_position' in apos:
        self._onfunc('on_position', self._ts, apos['symbol'], apos['net_position'])

  def _on_cancel_submit(self, cancel_submit):
    order_id = cancel_submit.get('proc_order_id', cancel_submit.get('internal_order_id'))
    symbol = cancel_submit['symbol']
    self._onfunc('on_cancel_submit', self._ts, symbol, order_id)

  def _on_cancel_accept(self, cancel_accept):
    order_id = cancel_accept.get('proc_order_id', cancel_accept.get('internal_order_id'))
    symbol = cancel_accept['symbol']
    self._onfunc('on_cancel_accept', self._ts, symbol, order_id)

  def _on_cancel_confirm(self, cancel_confirm):
    order_id = cancel_confirm.get('proc_order_id', cancel_confirm.get('internal_order_id'))
    symbol = cancel_confirm['symbol']
    self._onfunc('on_cancel_confirm', self._ts, symbol, order_id)

  def _on_cancel_error(self, cancel_error):
    symbol = cancel_error['symbol']
    order_id = cancel_error.get('proc_order_id', cancel_error.get('internal_order_id'))
    self._onfunc('on_cancel_error',
                 self._ts,
                 symbol,
                 order_id,
                 cancel_error.get('error_code'),
                 cancel_error.get('exchange_error_code'),
                 cancel_error.get('exchange_error_msg'))

  def _on_cancel_reject(self, cancel_reject):
    symbol = cancel_reject['symbol']
    order_id = cancel_reject.get('proc_order_id', cancel_reject.get('internal_order_id'))
    self._onfunc('on_cancel_reject',
                 self._ts,
                 symbol,
                 order_id,
                 cancel_reject.get('error_code'),
                 cancel_reject.get('exchange_error_code'),
                 cancel_reject.get('exchange_error_msg'))

  def _on_config(self, config):
    self._onfunc('on_config', self._ts, config['symbol'],
                  config['lot_size'], config['min_pos'], config['max_pos'])

  def on_log(self, timestamp, log):
    if timestamp < self._ts:
      return
    self._ts = timestamp
    pb = StrategyLog()
    pb.ParseFromString(log)
    data = json_format.MessageToDict(pb, preserving_proto_field_name=True)
    if data['type'] == 'RESERVE':
      if 'exchange' in data['reserve']:
        self._on_reserve(data['reserve']['each_balance'])
    elif data['type'] == 'AGG_EXEUCTOR_CONFIG':
      self._on_config(data['agg_executor_config'])
    elif data['type'] == 'OG_LOG':
      og_log = data['og_log']

      if og_log['type'] == 'BALANCE':
        self._on_balance(og_log['balance'].get('each_balance', []))
      elif og_log['type'] == 'POSITION':
        if 'each_position' in og_log['position']:
          self._on_position(og_log['position']['each_position'])
      elif og_log['type'] == 'ORDER_EVENT':
        event = og_log['event']
        if event['type'] == 'ORDER_SUBMITTED':
          self._on_submit(event)
        elif event['type'] == 'ORDER_REJECTED':
          self._on_submit_reject(event)
        elif event['type'] == 'ORDER_ACCEPTED':
          self._on_order_accept(event)
        elif event['type'] == 'ORDER_FILLED':
          self._on_fill(event)
        elif event['type'] == 'CANCEL_SUBMITTED':
          self._on_cancel_submit(event)
        elif event['type'] == 'CANCEL_ACCEPTED':
          self._on_cancel_accept(event)
        elif event['type'] == 'CANCEL_CONFIRMED':
          self._on_cancel_confirm(event)
        elif event['type'] == 'CANCEL_ERROR':
          self._on_cancel_error(event)
        elif event['type'] == 'CANCEL_REJECTED':
          self._on_cancel_reject(event)


class _OnLogPB(_OnLog):
  def on_symbol_info(self, symbol_info):
    self.exchange = symbol_info.exchange
    self.symbol = symbol_info.product_name
    # TODO(jingyuan): fix timestamp
    self._onfunc('on_reserve', 0, self.symbol, symbol_info.reserve)

  def _on_submit(self, submit):
    order_id = submit.get('order_id')
    order_side = OrderSide.BUY if submit['sign'] == 1 else OrderSide.SELL
    order_price = submit['order_price']
    order_qty = submit['order_qty']
    self._onfunc('on_submit', self._ts, self.symbol, order_id, order_side, order_price, order_qty)

  def _on_order_accept(self, ack):
    order_id = ack.get('order_id')
    self._onfunc('on_order_accept', self._ts, self.symbol, order_id)

  def _on_fill(self, fill):
    order_id = fill.get('order_id')
    order_side = OrderSide.BUY if fill['sign'] == 1 else OrderSide.SELL
    fill_price = fill['fill_price']
    fill_qty = fill['fill_qty']
    maker = None
    fill_type = fill.get('fill_type')
    if fill_type is not None:
      maker = 'MAKER' in fill_type
    self._onfunc('on_fill',
                 self._ts,
                 self.symbol,
                 order_id,
                 order_side,
                 fill_price,
                 fill_qty,
                 maker)

  def _on_cancel_submit(self, cancel_submit):
    order_id = cancel_submit.get('order_id')
    self._onfunc('on_cancel_submit', self._ts, self.symbol, order_id)

  def _on_cancel_confirm(self, cancel_confirm):
    order_id = cancel_confirm.get('order_id')
    self._onfunc('on_cancel_confirm', self._ts, self.symbol, order_id)

  def _on_cancel_error(self, cancel_error):
    order_id = cancel_error.get('order_id')
    order_side = OrderSide.BUY if cancel_error['sign'] == 1 else OrderSide.SELL
    order_price = cancel_error['order_price']
    # no order qty info
    self._onfunc('on_cancel_error', self._ts, self.symbol, order_id, order_side, order_price, 0)

  def on_log(self, timestamp, log):
    if timestamp < self._ts:
      return
    self._ts = timestamp

    event = json_format.MessageToDict(log, preserving_proto_field_name=True)

    if event['type'] == 'ORDER_SUBMITTED':
      self._on_submit(event)
    if event['type'] == 'ORDER_ACCEPTED':
      self._on_order_accept(event)
    elif event['type'] == 'ORDER_FILLED':
      self._on_fill(event)
    elif event['type'] == 'CANCEL_SUBMITTED':
      self._on_cancel_submit(event)
    elif event['type'] == 'CANCEL_CONFIRMED':
      self._on_cancel_confirm(event)
    elif event['type'] == 'CANCEL_ERROR':
      self._on_cancel_error(event)


def run_from_strat_log(machine,
                       strategy_name,
                       start_time,
                       on_log_obj,
                       end_time=None,
                       root_dir='/remote/iosg/strat-1/buckets/log.raw.coin/live/strat_proto_log'):
  if end_time is None:
    end_time = start_time + datetime.timedelta(days=1)
  on_log = _OnLog(on_log_obj)
  run_from_strat_log_archive(on_log_callback=on_log.on_log,
                             start_time=start_time,
                             end_time=end_time,
                             root_dir=root_dir,
                             machine=machine,
                             strategy_name=strategy_name)


def run_from_sim_oe_pblog(oe_pathes, on_log_obj, start_ts, end_ts, symbol_filter=None):
  for oe_path in glob.glob(oe_pathes):
    with open(oe_path, 'rb') as fread:
      responses = sspec.OrderResponsesProto()
      responses.ParseFromString(fread.read())
      on_sim_log = _OnLogPB(on_log_obj)

      for response in responses.order_responses:
        if symbol_filter is None or symbol_filter in response.symbol_info.product_name:
          on_sim_log.on_symbol_info(response.symbol_info)
          if len(response.responses) > 0 \
            and (response.responses[0].time <= start_ts <= response.responses[-1].time
                 or response.responses[0].time <= end_ts <= response.responses[-1].time):
            for resp in response.responses:
              on_sim_log.on_log(resp.time, resp)
